package main

import (
	"bufio"
	"fmt"
	"os"
	"strconv"
	"strings"
)

type (
	Token    any
	Number   float64
	String   string
	Boolean  bool
	Symbol   string
	List     []Token
	Function struct {
		body   any
		env    *Environment
		params []Symbol
	}
	Closure struct {
		fn  Function
		env *Environment
	}
)

type Environment struct {
	vars  map[string]any
	outer *Environment
}

func NewEnvironment(outer *Environment) *Environment {
	return &Environment{
		vars:  make(map[string]any),
		outer: outer,
	}
}

func (e *Environment) find(name string) *Environment {
	if _, exists := e.vars[name]; exists {
		return e
	}
	if e.outer != nil {
		return e.outer.find(name)
	}
	return nil
}

func tokenize(program string) []Token {
	tokens := make([]Token, 0)
	reader := strings.NewReader(program)

	for {
		ch, _, err := reader.ReadRune()
		if err != nil {
			break
		}

		switch ch {
		case ' ', '\t', '\n', '\r':
			continue
		case '(', ')', '\'': // Added '\'' to handle quoted lists
			tokens = append(tokens, Symbol(string(ch)))
		case '"':
			var str strings.Builder
			for {
				ch, _, err = reader.ReadRune()
				if err != nil || ch == '"' {
					break
				}
				if ch == '\\' {
					ch, _, _ = reader.ReadRune()
					if ch == 'n' {
						ch = '\n'
					}
				}
				str.WriteRune(ch)
			}
			tokens = append(tokens, String(str.String()))
		case ';': // Check for comment start
			for {
				ch, _, err = reader.ReadRune()
				if err != nil || ch == '\n' {
					break
				}
			}
		default:
			var token strings.Builder
			token.WriteRune(ch)
			for {
				ch, _, err = reader.ReadRune()
				if err != nil {
					break
				}
				if strings.ContainsRune(" ()\t\n\r\"'", ch) { // Modified to include '\''
					reader.UnreadRune()
					break
				}
				token.WriteRune(ch)
			}

			tokenStr := token.String()
			switch tokenStr {
			case "t":
				tokens = append(tokens, Boolean(true))
			case "nil":
				tokens = append(tokens, Boolean(false))
			default:
				if num, err := strconv.ParseFloat(tokenStr, 64); err == nil {
					tokens = append(tokens, Number(num))
				} else {
					tokens = append(tokens, Symbol(tokenStr))
				}
			}
		}
	}
	return tokens
}

func parse(tokens []Token) (any, []Token) {
	if len(tokens) == 0 {
		return nil, tokens
	}

	token := tokens[0]
	tokens = tokens[1:]

	switch sym := token.(type) {
	case Symbol:
		switch string(sym) {
		case "(":
			list := make(List, 0)
			for len(tokens) > 0 {
				if sym, ok := tokens[0].(Symbol); ok && sym == ")" {
					return list, tokens[1:]
				}
				exp, remaining := parse(tokens)
				list = append(list, exp)
				tokens = remaining
			}
		case "'": // Handle quoted lists
			exp, remaining := parse(tokens)
			return List{Symbol("quote"), exp}, remaining
		default:
			return token, tokens
		}
	default:
		return token, tokens
	}
	return token, tokens
}

func eval(exp any, env *Environment) any {
	switch exp := exp.(type) {
	case Number:
		return exp
	case String:
		return exp
	case Boolean:
		return exp
	case Symbol:
		if e := env.find(string(exp)); e != nil {
			return e.vars[string(exp)]
		}
		return nil
	case List:
		if len(exp) == 0 {
			return nil
		}

		switch first := exp[0].(type) {
		case Symbol:
			switch string(first) {
			case "if":
				condition := eval(exp[1], env)
				switch v := condition.(type) {
				case Boolean:
					if v {
						return eval(exp[2], env)
					}
					return eval(exp[3], env)
				case Number:
					if v != 0 {
						return eval(exp[2], env)
					}
					return eval(exp[3], env)
				default:
					if v != nil {
						return eval(exp[2], env)
					}
					return eval(exp[3], env)
				}
			case "define":
				symbol := exp[1].(Symbol)
				value := eval(exp[2], env)
				env.vars[string(symbol)] = value
				return nil
			case "lambda":
				params := make([]Symbol, 0)
				for _, param := range exp[1].(List) {
					params = append(params, param.(Symbol))
				}
				return Closure{fn: Function{params: params, body: exp[2]}, env: env}
			case "+":
				args := make([]interface{}, 0)
				for _, arg := range exp[1:] {
					args = append(args, eval(arg, env))
				}

				hasString := false
				for _, arg := range args {
					if _, ok := arg.(String); ok {
						hasString = true
						break
					}
				}

				if hasString {
					var result strings.Builder
					for _, arg := range args {
						switch v := arg.(type) {
						case String:
							result.WriteString(string(v))
						case Number:
							result.WriteString(fmt.Sprintf("%g", float64(v)))
						case Boolean:
							result.WriteString(fmt.Sprintf("%v", bool(v)))
						default:
							result.WriteString(formatResult(v))
						}
					}
					return String(result.String())
				} else {
					sum := 0.0
					for _, arg := range args {
						if num, ok := arg.(Number); ok {
							sum += float64(num)
						}
					}
					return Number(sum)
				}
			case "-", "*", "/":
				args := make([]float64, 0)
				for _, arg := range exp[1:] {
					if num, ok := eval(arg, env).(Number); ok {
						args = append(args, float64(num))
					}
				}

				switch string(first) {
				case "-":
					result := args[0]
					for _, arg := range args[1:] {
						result -= arg
					}
					return Number(result)
				case "*":
					result := 1.0
					for _, arg := range args {
						result *= arg
					}
					return Number(result)
				case "/":
					result := args[0]
					for _, arg := range args[1:] {
						result /= arg
					}
					return Number(result)
				}
			case "=":
				if len(exp) != 3 {
					return nil
				}
				left := eval(exp[1], env)
				right := eval(exp[2], env)
				switch l := left.(type) {
				case Number:
					if r, ok := right.(Number); ok {
						return Boolean(l == r)
					}
				case String:
					if r, ok := right.(String); ok {
						return Boolean(l == r)
					}
				case Boolean:
					if r, ok := right.(Boolean); ok {
						return Boolean(l == r)
					}
				}
				return Boolean(false)
			case "<":
				if len(exp) != 3 {
					return nil
				}
				left := eval(exp[1], env)
				right := eval(exp[2], env)
				if l, ok := left.(Number); ok {
					if r, ok := right.(Number); ok {
						return Boolean(l < r)
					}
				}
				return Boolean(false)
			case ">":
				if len(exp) != 3 {
					return nil
				}
				left := eval(exp[1], env)
				right := eval(exp[2], env)
				if l, ok := left.(Number); ok {
					if r, ok := right.(Number); ok {
						return Boolean(l > r)
					}
				}
				return Boolean(false)
			case "<=":
				if len(exp) != 3 {
					return nil
				}
				left := eval(exp[1], env)
				right := eval(exp[2], env)
				if l, ok := left.(Number); ok {
					if r, ok := right.(Number); ok {
						return Boolean(l <= r)
					}
				}
				return Boolean(false)
			case ">=":
				if len(exp) != 3 {
					return nil
				}
				left := eval(exp[1], env)
				right := eval(exp[2], env)
				if l, ok := left.(Number); ok {
					if r, ok := right.(Number); ok {
						return Boolean(l >= r)
					}
				}
				return Boolean(false)
			case "not":
				if len(exp) != 2 {
					return nil
				}
				val := eval(exp[1], env)
				if b, ok := val.(Boolean); ok {
					return Boolean(!b)
				}
				return Boolean(false)
			case "cond":
				for i := 1; i < len(exp); i++ {
					clause := exp[i].(List)
					test := clause[0]
					result := clause[1]
					if eval(test, env).(Boolean) {
						return eval(result, env)
					}
				}
				return nil
			case "append":
				if len(exp) < 3 {
					return nil
				}
				firstList := eval(exp[1], env).(List)
				for i := 2; i < len(exp); i++ {
					nextItem := eval(exp[i], env)
					if nextList, ok := nextItem.(List); ok {
						firstList = append(firstList, nextList...)
					} else {
						firstList = append(firstList, nextItem)
					}
				}
				return firstList
			case "remove":
				if len(exp) != 3 {
					return nil
				}
				itemToRemove := eval(exp[1], env)
				listToModify := eval(exp[2], env).(List)
				newList := make(List, 0)
				for _, item := range listToModify {
					if !equal(item, itemToRemove) {
						newList = append(newList, item)
					}
				}
				return newList
			case "remove-if":
				if len(exp) != 3 {
					return nil
				}
				predicate := exp[1].(Closure)
				listToModify := eval(exp[2], env).(List)
				newList := make(List, 0)
				for _, item := range listToModify {
					if !eval(predicate.fn.body, NewEnvironmentWithBinding(predicate.env, predicate.fn.params[0], item)).(Boolean) {
						newList = append(newList, item)
					}
				}
				return newList
			case "remove-if-not":
				if len(exp) != 3 {
					return nil
				}
				predicate := exp[1].(Closure)
				listToModify := eval(exp[2], env).(List)
				newList := make(List, 0)
				for _, item := range listToModify {
					if eval(predicate.fn.body, NewEnvironmentWithBinding(predicate.env, predicate.fn.params[0], item)).(Boolean) {
						newList = append(newList, item)
					}
				}
				return newList
			case "first":
				if len(exp) != 2 {
					return nil
				}
				list := eval(exp[1], env).(List)
				if len(list) > 0 {
					return list[0]
				}
				return nil
			case "rest":
				if len(exp) != 2 {
					return nil
				}
				list := eval(exp[1], env).(List)
				if len(list) > 0 {
					return list[1:]
				}
				return nil
			case "nth":
				if len(exp) != 3 {
					return nil
				}
				n := int(eval(exp[1], env).(Number))
				list := eval(exp[2], env).(List)
				if n >= 0 && n < len(list) {
					return list[n]
				}
				return nil
			case "setf":
				if len(exp) != 3 {
					return nil
				}
				target := exp[1]
				value := eval(exp[2], env)
				setValue(target, value, env)
				return value
			case "defun":
				if len(exp) < 3 {
					return nil
				}
				name := exp[1].(Symbol)
				lambdaExp := List{Symbol("lambda"), exp[2], exp[3]}
				function := eval(lambdaExp, env).(Closure)
				env.vars[string(name)] = function
				return nil
			case "list":
				return exp[1:]
			case "cons":
				if len(exp) != 3 {
					return nil
				}
				head := eval(exp[1], env)
				tail := eval(exp[2], env).(List)
				return append([]Token{head}, tail...)
			case "quote":
				return exp[1]
			case "let":
				if len(exp) < 3 {
					return nil
				}
				bindings := exp[1].(List)
				body := exp[2]
				letEnv := NewEnvironment(env)
				for _, binding := range bindings {
					bindingList := binding.(List)
					varName := bindingList[0].(Symbol)
					varValue := eval(bindingList[1], letEnv)
					letEnv.vars[string(varName)] = varValue
				}
				return eval(body, letEnv)
			case "null":
				if len(exp) != 2 {
					return nil
				}
				list := eval(exp[1], env)
				if list == nil || (len(list.(List)) == 0) {
					return Boolean(true)
				}
				return Boolean(false)
			default:
				fn := eval(first, env)
				if closure, ok := fn.(Closure); ok {
					newEnv := NewEnvironment(closure.env)
					for i, param := range closure.fn.params {
						newEnv.vars[string(param)] = eval(exp[i+1], env)
					}
					return eval(closure.fn.body, newEnv)
				}
			}
		}
	}
	return nil
}

func equal(a, b Token) bool {
	switch a := a.(type) {
	case Number:
		if b, ok := b.(Number); ok {
			return a == b
		}
	case String:
		if b, ok := b.(String); ok {
			return a == b
		}
	case Boolean:
		if b, ok := b.(Boolean); ok {
			return a == b
		}
	case Symbol:
		if b, ok := b.(Symbol); ok {
			return a == b
		}
	case List:
		if b, ok := b.(List); ok {
			if len(a) != len(b) {
				return false
			}
			for i := range a {
				if !equal(a[i], b[i]) {
					return false
				}
			}
			return true
		}
	}
	return false
}

func setValue(target any, value any, env *Environment) {
	switch target := target.(type) {
	case Symbol:
		env.find(string(target)).vars[string(target)] = value
	case List:
		if len(target) == 2 {
			first := target[0]
			if firstSym, ok := first.(Symbol); ok {
				if firstSym == Symbol("nth") {
					index := eval(target[1], env).(Number)
					listVar := target[2].(Symbol)
					list := env.find(string(listVar)).vars[string(listVar)].(List)
					if index >= 0 && index < Number(len(list)) {
						list[int(index)] = value
						env.find(string(listVar)).vars[string(listVar)] = list
					}
				}
			}
		}
	}
}

func formatResult(result any) string {
	switch v := result.(type) {
	case Number:
		return fmt.Sprintf("%g", float64(v))
	case String:
		return string(v)
	case Boolean:
		if v {
			return "t"
		}
		return "nil"
	case Symbol:
		return string(v)
	case List:
		var sb strings.Builder
		sb.WriteString("(")
		for i, item := range v {
			sb.WriteString(formatResult(item))
			if i < len(v)-1 {
				sb.WriteString(" ")
			}
		}
		sb.WriteString(")")
		return sb.String()
	case Closure:
		return "#<CLOSURE>"
	default:
		return fmt.Sprintf("%v", v)
	}
}

func NewEnvironmentWithBinding(outer *Environment, key Symbol, value any) *Environment {
	env := NewEnvironment(outer)
	env.vars[string(key)] = value
	return env
}

func main() {
	if len(os.Args) != 2 {
		fmt.Println("Usage: mylisp <filename.lisp>")
		os.Exit(1)
	}
	file, err := os.Open(os.Args[1])
	if err != nil {
		fmt.Printf("Error opening file: %v\n", err)
		return
	}
	defer file.Close()
	globalEnv := NewEnvironment(nil)
	scanner := bufio.NewScanner(file)
	var program strings.Builder
	for scanner.Scan() {
		line := scanner.Text()
		program.WriteString(line)
		program.WriteString("\n")
	}
	if err := scanner.Err(); err != nil {
		fmt.Printf("Error reading file: %v\n", err)
		return
	}

	tokens := tokenize(program.String())
	for len(tokens) > 0 {
		ast, remaining := parse(tokens)
		if ast != nil {
			result := eval(ast, globalEnv)
			if result != nil {
				fmt.Println(formatResult(result))
			}
		}
		tokens = remaining
	}
}
